"""
postgresql 客户端，单客户端，连接池
"""

import psycopg2
from contextlib import contextmanager
from playhouse.shortcuts import ReconnectMixin
from playhouse.pool import PooledPostgresqlDatabase
from psycopg2 import extras
import time
import logging


class CursorWithLog(extras.DictCursorBase):
    def execute(self, query, vars=None):
        """Execute a query
                        带日志输出的SQL执行。改造了Cursor类，自动输出日志，方便Debug和查问题。
                        :param str query: Query to execute.

                        :param vars: parameters used with query. (optional)
                        :type vars: tuple, list or dict

                        :return: Number of affected rows
                        :rtype: int

                        If args is a list or tuple, %s can be used as a placeholder in the query.
                        If args is a dict, %(name)s can be used as a placeholder in the query.
                        """
        start_time = time.time()
        try:
            result = super().execute(query, vars)
        except Exception as e:
            raise Exception(f"【 {time.asctime()} 】\nFail to execute SQL: 【 {query} 】 \nDATA:  【 {vars} 】")
        end_time = time.time()
        cost_time = (end_time - start_time) * 1000
        msg = f"【 {time.asctime()} 】\nSuccess to execute SQL with {cost_time} ms \nSQL: 【 {query} 】 \nDATA:  【 {vars} 】"
        logging.warning(msg)
        return result


class DictCursor(extras.DictCursor, CursorWithLog):
    pass


class RealDictCursor(extras.RealDictCursor, CursorWithLog):
    pass


class NamedTupleCursor(extras.NamedTupleCursor, CursorWithLog):
    pass


class PgClient():
    """
    Create a new database connection.

    The connection parameters can be specified as a string:

        conn = psycopg2.connect("dbname=test user=postgres password=secret")

    or using a set of keyword arguments:

        conn = psycopg2.connect(database="test", user="postgres", password="secret")

    Or as a mix of both. The basic connection parameters are:

    - *dbname*: the database name
    - *database*: the database name (only as keyword argument)
    - *user*: user name used to authenticate
    - *password*: password used to authenticate
    - *host*: database host address (defaults to UNIX socket if not provided)
    - *port*: connection port number (defaults to 5432 if not provided)

    Using the *connection_factory* parameter a different class or connections
    factory can be specified. It should be a callable object taking a dsn
    argument.

    Using the *cursor_factory* parameter, a new default cursor factory will be
    used by cursor().

    Using *async*=True an asynchronous connection will be created. *async_* is
    a valid alias (for Python versions where ``async`` is a keyword).

    Any other keyword parameter will be passed to the underlying client
    library: the list of supported parameters depends on the library version.
    """
    pg_conf = {}

    def __init__(self, **pg_conf):
        if pg_conf:
            self.pgConfig = pg_conf
        else:
            self.pgConfig = self.pg_conf

        self.__connection__ = None

    @property
    def connection(self):
        return self.__connection__

    def connect(self):
        self.__connection__ = psycopg2.connect(**self.pgConfig)

    def reconnect(self):
        self.__connection__ = psycopg2.connect(**self.pgConfig)

    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            self.__connection__.close()
        except Exception as e:
            pass


@contextmanager
def PgClientSession(pg: PgClient):
    """
    pg原生客户端session
    :param pg:
    :return:
    """
    _Exception = False
    cursor = None
    connection = None
    try:
        pg.connect()
        connection = pg.connection
        cursor = pg.connection.cursor()
        yield cursor
        pg.connection.commit()
    except Exception as e:
        _Exception = e
        pg.connection.rollback()
    finally:
        try:
            cursor.close()
            connection.close()
        except:
            pass
    if _Exception:
        raise Exception(_Exception)


def PeeWeePgPool(**pg_conf) -> PooledPostgresqlDatabase:
    """
    Create a new database connection pool.


    - *dbname*: the database name
    - *database*: the database name (only as keyword argument)
    - *user*: user name used to authenticate
    - *password*: password used to authenticate
    - *host*: database host address (defaults to UNIX socket if not provided)
    - *port*: connection port number (defaults to 5432 if not provided)
    - *max_connections*: connection number
    - *timeout*: timeout

    Using the *connection_factory* parameter a different class or connections
    factory can be specified. It should be a callable object taking a dsn
    argument.

    Using the *cursor_factory* parameter, a new default cursor factory will be
    used by cursor().

    Using *async*=True an asynchronous connection will be created. *async_* is
    a valid alias (for Python versions where ``async`` is a keyword).

    Any other keyword parameter will be passed to the underlying client
    library: the list of supported parameters depends on the library version.
    """

    class ReconnectPooledPostgresqlDatabase(ReconnectMixin, PooledPostgresqlDatabase):
        def execute_sql(self, sql, params=None, commit=object()):
            try:
                return super(ReconnectMixin, self).execute_sql(sql, params, commit)
            except Exception as exc:
                if not self.is_closed():
                    self.close()
                    self.connect()

                return super(ReconnectMixin, self).execute_sql(sql, params, commit)

    pool = ReconnectPooledPostgresqlDatabase(**pg_conf)
    return pool


@contextmanager
def PeeWeePgPoolSession(pool: PeeWeePgPool) -> PeeWeePgPool:
    """
    基于peewe封装的pg连接池，支持自动提交，自动回滚
    :param pool:
    :return:
    """
    _Exception = None
    try:
        with pool.atomic():
            yield pool.cursor()
    finally:
        pool.close()
    if _Exception:
        raise _Exception


def test():
    postgresql_host = "172.31.3.188"
    postgresql_port = "55433"
    postgresql_database = "ai_bigdata"
    postgresql_user = "postgres"
    postgresql_password = "abc123"
    p = PeeWeePgPool(database=postgresql_database, user=postgresql_user,
                     password=postgresql_password, host=postgresql_host, cursor_factory=RealDictCursor,
                     port=postgresql_port, max_connections=3, timeout=10)

    def thr():
        with PeeWeePgPoolSession(p) as session:
            # a = session.execute("insert into age  (age) values (5)")
            session.execute("select 1")
            print(session.fetchmany())
            # raise # 测试事务

    from threading import Thread
    for i in range(2):
        Thread(target=thr).start()

    time.sleep(100)


if __name__ == '__main__':
    test()
